"""Contains a variant of the densenet model definition."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

slim = tf.contrib.slim

#返回一个正态分布的初始化器
def trunc_normal(stddev): return tf.truncated_normal_initializer(stddev=stddev)


#定义每一层layer包含的操作
def bn_act_conv_drp(current, num_outputs, kernel_size, scope='block'):
    current = slim.batch_norm(current, scope=scope + '_bn')
    current = tf.nn.relu(current)
    current = slim.conv2d(current, num_outputs, kernel_size, scope=scope + '_conv')
    current = slim.dropout(current, scope=scope + '_dropout')
    return current

#block的构造函数
def block(net, layers, growth, scope='block'):
    for idx in range(layers):                                                      #每个block分为不同层layer
        bottleneck = bn_act_conv_drp(net, 4 * growth, [1, 1],  
                                     scope=scope + '_conv1x1' + str(idx))           #瓶颈层，用【1,1】卷积核降维，减少channel为 4*growth.每一个layer。
        tmp = bn_act_conv_drp(bottleneck, growth, [3, 3],
                              scope=scope + '_conv3x3' + str(idx))                  #降维之后再用 【3,3】卷积，channel为growth。之前所有layer的输出。
        net = tf.concat(axis=3, values=[net, tmp])                                  #拼接layer。稠密连接
    return net


def densenet(images, num_classes=1001, is_training=False,
             dropout_keep_prob=0.8,
             scope='densenet'):
    """Creates a variant of the densenet model.

      images: A batch of `Tensors` of size [batch_size, height, width, channels].
      num_classes: the number of classes in the dataset.
      is_training: specifies whether or not we're currently training the model.
        This variable will determine the behaviour of the dropout layer.
      dropout_keep_prob: the percentage of activation values that are retained.
      prediction_fn: a function to get predictions out of logits.
      scope: Optional variable_scope.

    Returns:
      logits: the pre-softmax activations, a tensor of size
        [batch_size, `num_classes`]
      end_points: a dictionary from components of the network to the corresponding
        activation.
    """
    growth = 24                 #增长率，每次循环增加的featureMap数目
    compression_rate = 0.5      #压缩率，变换层降低channel

    #定义函数，返回变换层压缩后的channel数
    def reduce_dim(input_feature):
        return int(int(input_feature.shape[-1]) * compression_rate)

    end_points = {}

    with tf.variable_scope(scope, 'DenseNet', [images, num_classes]):
        with slim.arg_scope(bn_drp_scope(is_training=is_training,
                                         keep_prob=dropout_keep_prob)) as ssc:
            pass
            ##########################
            # Put your code here.
            #根据论文的Tabel1，实现DenseNet结构
            
            with tf.variable_scope('pre_processing') :                          #预处理层
                net=slim.conv2d(images, 2*growth, [7,7],stride=2,scope= 'pre_con2d')
                net=slim.max_pool2d(net, [3, 3], stride=2, scope='pre_pool')
                end_points['pre_processing']=net
                
            
            with tf.variable_scope('block1') :                                  #block1
                net=block(net, 6, growth, scope='block1')
                end_points['block1']=net
                
            
            with tf.variable_scope('transition1') :                            #变换层1
                net=bn_act_conv_drp(net, reduce_dim(net), [1,1], scope='trans1_con2d')
                net=slim.max_pool2d(net, [2, 2], stride=2, scope='trans1_pool')
                end_points['transition1']=net
                
            
            with tf.variable_scope('block2') :                                 #block2
                net=block(net, 12, growth, scope='block2')
                end_points['block2']=net
                
            
            with tf.variable_scope('transition2') :                            #变换层2
                net=bn_act_conv_drp(net, reduce_dim(net), [1,1], scope='trans2_con2d')
                net=slim.max_pool2d(net, [2, 2], stride=2, scope='trans2_pool')
                end_points['transition2']=net
                
                
            with tf.variable_scope('block3') :                                #block3
                net=block(net, 24, growth, scope='block3')
                end_points['block3']=net
                
            
            with tf.variable_scope('transition3') :                           #变换层3
                net=bn_act_conv_drp(net, reduce_dim(net), [1,1], scope='trans3_con2d')
                net=slim.max_pool2d(net, [2, 2], stride=2, scope='trans3_pool')
                end_points['transition3']=net
                
            
            with tf.variable_scope('block4') :                               #block4
                net=block(net, 16, growth, scope='block4')
                end_points['block4']=net
                
            
            with tf.variable_scope('classify_layer') :                       #分类层
                net=slim.max_pool2d(net, net.shape[1:3], scope='global_average')
                net=slim.flatten(net, scope='flatten')
                logits =slim.fully_connected(net, num_classes, activation_fn=None, scope='logits')
                net=tf.nn.softmax(logits, name='predictions')
                end_points['classify_layer']=net
            
            ##########################

    return logits, end_points


#定义使用batch_norm与dropout函数时，默认的共同参数
def bn_drp_scope(is_training=True, keep_prob=0.8):
    keep_prob = keep_prob if is_training else 1
    with slim.arg_scope(
        [slim.batch_norm],
            scale=True, is_training=is_training, updates_collections=None):
        with slim.arg_scope(
            [slim.dropout],
                is_training=is_training, keep_prob=keep_prob) as bsc:
            return bsc



#定义使用conv2d函数时，默认的共同参数
def densenet_arg_scope(weight_decay=0.004):
    """Defines the default densenet argument scope.

    Args:
      weight_decay: The weight decay to use for regularizing the model.

    Returns:
      An `arg_scope` to use for the inception v3 model.
    """
    with slim.arg_scope(
        [slim.conv2d],
        weights_initializer=tf.contrib.layers.variance_scaling_initializer(
            factor=2.0, mode='FAN_IN', uniform=False),
        activation_fn=None, biases_initializer=None, padding='same',
            stride=1) as sc:
        return sc


densenet.default_image_size = 224
